> ## Documentation Index > Fetch the complete documentation index at: https://mintlify.com/mlfoundations/open_clip/llms.txt > Use this file to discover all available pages before exploring further. # CoCa Training > Train CoCa (Contrastive Captioner) models for image captioning and contrastive learning ## What is CoCa? CoCa (Contrastive Captioner) is an extension of CLIP that combines: 1. **Contrastive Learning**: Standard CLIP image-text matching 2. **Generative Captioning**: Auto-regressive caption generation This dual objective enables CoCa models to: * Perform zero-shot image classification (like CLIP) * Generate natural language captions for images * Achieve better representations through the combined training signal **Paper**: [CoCa: Contrastive Captioners are Image-Text Foundation Models](https://arxiv.org/abs/2205.01917) ## CoCa Architecture CoCa adds a multimodal text decoder on top of the standard CLIP architecture: ``` [Image] → Image Encoder → Image Features →┬─ Contrastive Loss │ [Text] → Text Encoder → Text Features ─┴─ Caption Loss │ └─ Multimodal Decoder → Generated Caption ``` **Key components:** * **Image Encoder**: Same as CLIP (ViT, ResNet, etc.) * **Unimodal Text Encoder**: Encodes text for contrastive learning * **Multimodal Text Decoder**: Cross-attends to image features to generate captions ## Available CoCa Models OpenCLIP provides several CoCa model configurations: ### Model Configs | Model | Image Encoder | Text Encoder | Multimodal Decoder | | ----------------------- | ------------- | ------------ | -------------------- | | `coca_base` | ViT-B/16 | Transformer | 12-layer Transformer | | `coca_ViT-B-32` | ViT-B/32 | Transformer | 12-layer Transformer | | `coca_ViT-L-14` | ViT-L/14 | Transformer | 12-layer Transformer | | `coca_roberta-ViT-B-32` | ViT-B/32 | RoBERTa | 12-layer Transformer | ### Multimodal Decoder Configuration Example configuration from `coca_ViT-B-32`: ```json theme={null} "multimodal_cfg": { "context_length": 76, "vocab_size": 49408, "width": 512, "heads": 8, "layers": 12, "latent_dim": 512, "attn_pooler_heads": 8 } ``` ## Training CoCa from Scratch ### Basic CoCa Training Train CoCa with both contrastive and captioning objectives: ```bash theme={null} python -m open_clip_train.main \ --model coca_ViT-L-14 \ --train-data "/data/train.tar" \ --train-num-samples 10000000 \ --dataset-type webdataset \ --batch-size 128 \ --precision amp \ --workers 8 \ --lr 1e-3 \ --wd 0.1 \ --warmup 10000 \ --epochs 32 \ --coca-contrastive-loss-weight 1.0 \ --coca-caption-loss-weight 2.0 \ --report-to wandb ``` **Loss weights:** * `--coca-contrastive-loss-weight 1.0`: Weight for CLIP contrastive loss * `--coca-caption-loss-weight 2.0`: Weight for caption generation loss ### Multi-GPU CoCa Training ```bash theme={null} torchrun --nproc_per_node 8 -m open_clip_train.main \ --model coca_ViT-L-14 \ --train-data "/data/laion400m/{00000..41455}.tar" \ --train-num-samples 400000000 \ --dataset-type webdataset \ --dataset-resampled \ --batch-size 128 \ --precision amp \ --grad-checkpointing \ --workers 8 \ --lr 5e-4 \ --wd 0.2 \ --warmup 10000 \ --epochs 32 \ --coca-contrastive-loss-weight 1.0 \ --coca-caption-loss-weight 2.0 \ --local-loss \ --gather-with-grad \ --report-to wandb ``` ## Fine-tuning CoCa ### Fine-tuning on MSCOCO Captions OpenCLIP provides a pretrained CoCa model that can be fine-tuned for captioning: ```bash theme={null} python -m open_clip_train.main \ --model coca_ViT-L-14 \ --pretrained laion2b_s13b_b90k \ --dataset-type csv \ --train-data "/data/mscoco/train2014.csv" \ --csv-img-key filepath \ --csv-caption-key title \ --csv-separator "\t" \ --warmup 1000 \ --batch-size 128 \ --lr 1e-5 \ --wd 0.1 \ --epochs 1 \ --workers 3 \ --coca-contrastive-loss-weight 0 \ --coca-caption-loss-weight 1 \ --report-to wandb \ --log-every-n-steps 100 ``` **Key changes for fine-tuning:** * `--pretrained laion2b_s13b_b90k`: Start from pretrained weights * `--lr 1e-5`: Lower learning rate for fine-tuning * `--epochs 1`: Fine-tune for fewer epochs * `--coca-contrastive-loss-weight 0`: Disable contrastive loss (captioning only) * `--coca-caption-loss-weight 1`: Only train the generative head ### Preparing MSCOCO Data Create a CSV file with image paths and captions using [CLIP\_benchmark](https://github.com/LAION-AI/CLIP_benchmark): ```python theme={null} from clip_benchmark.datasets.builder import build_dataset import pandas as pd import os root_path = "path/to/data/dir" # Set this to your data directory # Download and load MSCOCO ds = build_dataset("mscoco_captions", root=root_path, split="train", task="captioning") coco = ds.coco imgs = coco.loadImgs(coco.getImgIds()) # Create CSV with all image-caption pairs future_df = {"filepath": [], "title": []} for img in imgs: caps = coco.imgToAnns[img["id"]] for cap in caps: future_df["filepath"].append(img["file_name"]) future_df["title"].append(cap["caption"]) # Save to CSV pd.DataFrame.from_dict(future_df).to_csv( os.path.join(root_path, "train2014.csv"), index=False, sep="\t" ) ``` This creates a tab-separated CSV: ```tsv theme={null} filepath\ttitle COCO_train2014_000000000009.jpg\tA person on a motorcycle on a dirt road COCO_train2014_000000000009.jpg\tA man riding a motorcycle down a dirt road ... ``` ## Generating Captions with CoCa ### Basic Caption Generation ```python theme={null} import open_clip import torch from PIL import Image # Load pretrained CoCa model model, _, transform = open_clip.create_model_and_transforms( model_name="coca_ViT-L-14", pretrained="mscoco_finetuned_laion2B-s13B-b90k" ) # Load and preprocess image im = Image.open("cat.jpg").convert("RGB") im = transform(im).unsqueeze(0) # Generate caption with torch.no_grad(), torch.cuda.amp.autocast(): generated = model.generate(im) # Decode and print caption = open_clip.decode(generated[0]).split("")[0].replace("", "") print(caption) # Output: "a cat sitting on a windowsill" ``` ### Batch Caption Generation ```python theme={null} import open_clip import torch from PIL import Image model, _, transform = open_clip.create_model_and_transforms( model_name="coca_ViT-L-14", pretrained="mscoco_finetuned_laion2B-s13B-b90k" ) model.eval() # Load multiple images images = [ Image.open("cat.jpg").convert("RGB"), Image.open("dog.jpg").convert("RGB"), Image.open("car.jpg").convert("RGB"), ] # Preprocess images_tensor = torch.stack([transform(im) for im in images]) # Generate captions with torch.no_grad(), torch.cuda.amp.autocast(): generated = model.generate(images_tensor) # Decode all captions for i, gen in enumerate(generated): caption = open_clip.decode(gen).split("")[0].replace("", "") print(f"Image {i}: {caption}") ``` ### Advanced Generation Options ```python theme={null} # Generate with custom parameters with torch.no_grad(), torch.cuda.amp.autocast(): generated = model.generate( im, seq_len=77, # Maximum sequence length temperature=1.0, # Sampling temperature top_p=0.9, # Nucleus sampling ) ``` ## CoCa vs CLIP ### When to Use CoCa ✅ **Use CoCa when:** * You need both contrastive and generative capabilities * Image captioning is important for your application * You want richer image-text representations * You have data with detailed captions ### When to Use CLIP ✅ **Use CLIP when:** * You only need contrastive learning (classification, retrieval) * Training speed is critical (CoCa is slower due to caption generation) * You have limited compute resources * Your captions are short or simple ### Training Time Comparison | Model | Architecture | Training Speed (relative) | Memory Usage (relative) | | ------------- | ---------------------- | ------------------------- | ----------------------- | | CLIP ViT-L/14 | Dual encoder | 1.0× | 1.0× | | CoCa ViT-L/14 | Dual encoder + decoder | 0.6× | 1.4× | CoCa is slower due to the autoregressive caption generation during training. ## Example Training Configurations ### Small-Scale CoCa Training ```bash theme={null} # CoCa ViT-B/32 on CC3M torchrun --nproc_per_node 4 -m open_clip_train.main \ --model coca_ViT-B-32 \ --train-data "/data/cc3m/train-{0000..0331}.tar" \ --train-num-samples 3000000 \ --dataset-type webdataset \ --batch-size 256 \ --precision amp \ --workers 4 \ --lr 1e-3 \ --wd 0.1 \ --warmup 5000 \ --epochs 30 \ --coca-contrastive-loss-weight 1.0 \ --coca-caption-loss-weight 2.0 \ --report-to tensorboard ``` ### Large-Scale CoCa Training ```bash theme={null} # CoCa ViT-L/14 on LAION-400M srun python -u src/open_clip_train/main.py \ --model coca_ViT-L-14 \ --train-data "/data/laion400m/{00000..41455}.tar" \ --train-num-samples 400000000 \ --dataset-type webdataset \ --dataset-resampled \ --batch-size 128 \ --precision amp \ --grad-checkpointing \ --workers 8 \ --lr 5e-4 \ --wd 0.2 \ --warmup 10000 \ --epochs 32 \ --coca-contrastive-loss-weight 1.0 \ --coca-caption-loss-weight 2.0 \ --local-loss \ --gather-with-grad \ --report-to wandb \ --remote-sync s3://bucket/coca-checkpoints ``` ### CoCa with RoBERTa Text Encoder ```bash theme={null} # Use RoBERTa for the unimodal text encoder python -m open_clip_train.main \ --model coca_roberta-ViT-B-32 \ --train-data "/data/train.tar" \ --batch-size 256 \ --precision amp \ --workers 4 \ --lr 5e-4 \ --wd 0.1 \ --warmup 2000 \ --epochs 10 \ --lock-text \ --lock-text-unlocked-layers 10 \ --coca-contrastive-loss-weight 1.0 \ --coca-caption-loss-weight 2.0 ``` ## Pretrained CoCa Models OpenCLIP provides pretrained CoCa models: ```python theme={null} import open_clip # List available CoCa models models = open_clip.list_pretrained() coca_models = [m for m in models if 'coca' in m[0].lower()] for model_name, pretrained in coca_models: print(f"{model_name}: {pretrained}") # Load pretrained CoCa model, _, preprocess = open_clip.create_model_and_transforms( 'coca_ViT-L-14', pretrained='mscoco_finetuned_laion2B-s13B-b90k' ) ``` **Available pretrained weights:** * `laion2b_s13b_b90k`: Pretrained on LAION-2B * `mscoco_finetuned_laion2B-s13B-b90k`: LAION-2B pretraining + MSCOCO fine-tuning ## Using CoCa for Multiple Tasks ### Image Classification (Zero-Shot) ```python theme={null} import open_clip import torch from PIL import Image model, _, preprocess = open_clip.create_model_and_transforms( 'coca_ViT-L-14', pretrained='mscoco_finetuned_laion2B-s13B-b90k' ) tokenizer = open_clip.get_tokenizer('coca_ViT-L-14') image = preprocess(Image.open("dog.jpg")).unsqueeze(0) text = tokenizer(["a dog", "a cat", "a car"]) with torch.no_grad(), torch.cuda.amp.autocast(): image_features = model.encode_image(image) text_features = model.encode_text(text) image_features /= image_features.norm(dim=-1, keepdim=True) text_features /= text_features.norm(dim=-1, keepdim=True) similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1) print("Probabilities:", similarity) ``` ### Image Captioning ```python theme={null} # Same model, different usage with torch.no_grad(), torch.cuda.amp.autocast(): generated = model.generate(image) caption = open_clip.decode(generated[0]).split("")[0].replace("", "") print("Caption:", caption) ``` ### Image-Text Retrieval ```python theme={null} # Encode multiple images and texts images = torch.stack([preprocess(Image.open(f"image_{i}.jpg")) for i in range(10)]) texts = tokenizer([f"caption {i}" for i in range(100)]) with torch.no_grad(), torch.cuda.amp.autocast(): image_features = model.encode_image(images) text_features = model.encode_text(texts) # Normalize image_features /= image_features.norm(dim=-1, keepdim=True) text_features /= text_features.norm(dim=-1, keepdim=True) # Compute similarity matrix similarity = image_features @ text_features.T # [10, 100] # Top-5 captions for each image top5 = similarity.topk(5, dim=-1) print("Top 5 matches:", top5.indices) ``` ## Tips for Training CoCa **Training tips:** * Start with contrastive-only training, then add caption loss gradually * Use higher weight for caption loss (2.0 vs 1.0 for contrastive) * Fine-tune on high-quality caption datasets (MSCOCO) for best generation * Use gradient checkpointing for memory efficiency with large models * Caption generation is slower - expect 40-60% of CLIP training speed **Common issues:** * CoCa requires more memory due to the multimodal decoder * Gradient accumulation is not compatible with distillation for CoCa * Caption quality depends heavily on training data quality * Very short captions may not benefit from the generative objective ## Credits CoCa implementation in OpenCLIP: * Initial implementation: [lucidrains](https://github.com/lucidrains/CoCa-pytorch) * Adaptation to OpenCLIP: [gpucce](https://github.com/gpucce) * Training: [iejMac](https://github.com/iejMac) ## Next Steps Learn about general CLIP training Fine-tune CoCa models on custom datasets Explore all CoCa training parameters Use pretrained CoCa for captioning and classification